Step 1: Data processing¶

In [ ]:
from google.colab import drive
drive.mount('/content/drive')
Mounted at /content/drive
In [ ]:
import zipfile
import os

zip_file_path = '/content/drive/MyDrive/archive.zip'

extract_to_path = '/content/traffic_light'

os.makedirs(extract_to_path, exist_ok=True)

with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
    zip_ref.extractall(extract_to_path)

print(f"zip to: {extract_to_path}")
zip to: /content/traffic_light
In [11]:
import os
import json
import shutil
import cv2
from sklearn.model_selection import train_test_split
from concurrent.futures import ThreadPoolExecutor
In [ ]:
def split_dataset(json_file, images_dir, output_dir, split_ratios=(0.7, 0.15, 0.15)):
    with open(json_file, 'r') as f:
        data = json.load(f)
    annotations = data['annotations']

    os.makedirs(output_dir, exist_ok=True)
    split_dirs = {
        "train": os.path.join(output_dir, "train"),
        "val": os.path.join(output_dir, "val"),
        "test": os.path.join(output_dir, "test"),
    }

    for split in split_dirs.values():
        os.makedirs(os.path.join(split, "images"), exist_ok=True)
        os.makedirs(os.path.join(split, "labels"), exist_ok=True)

    grouped_annotations = {}
    for annotation in annotations:
        filename = annotation['filename'].replace("\\", "/")
        grouped_annotations.setdefault(filename, []).append(annotation)

    filenames = list(grouped_annotations.keys())
    train_files, temp_files = train_test_split(filenames, test_size=(1 - split_ratios[0]), random_state=42)
    val_files, test_files = train_test_split(temp_files, test_size=split_ratios[2] / (split_ratios[1] + split_ratios[2]), random_state=42)
    splits = {"train": train_files, "val": val_files, "test": test_files}

    def process_file(filename, annotations, split_dir):
        img_name = os.path.basename(filename)
        label_name = os.path.splitext(img_name)[0] + ".txt"

        src_image = os.path.join(images_dir, filename)
        dst_image = os.path.join(split_dir, "images", img_name)
        if not os.path.exists(src_image):
            print(f"Warning: {src_image} does not exist.")
            return

        shutil.copyfile(src_image, dst_image)

        img = cv2.imread(src_image)
        if img is None:
            print(f"Warning: Could not read {src_image}. Skipping.")
            return
        img_height, img_width = img.shape[:2]

        label_path = os.path.join(split_dir, "labels", label_name)
        with open(label_path, 'w') as f:
            for annotation in annotations:
                bbox = annotation['bndbox']
                color = annotation['inbox'][0]['color'] if annotation['inbox'] else 'red'
                class_id = 0 if color == 'red' else 1

                xmin, ymin, xmax, ymax = bbox['xmin'], bbox['ymin'], bbox['xmax'], bbox['ymax']
                xmin, xmax = max(0, xmin), min(img_width - 1, xmax)
                ymin, ymax = max(0, ymin), min(img_height - 1, ymax)

                x_center = (xmin + xmax) / 2 / img_width
                y_center = (ymin + ymax) / 2 / img_height
                width = (xmax - xmin) / img_width
                height = (ymax - ymin) / img_height

                if 0 <= x_center <= 1 and 0 <= y_center <= 1 and 0 <= width <= 1 and 0 <= height <= 1:
                    f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")

    def save_split(files, split_name):
        split_dir = split_dirs[split_name]
        with ThreadPoolExecutor() as executor:
            futures = [
                executor.submit(process_file, filename, grouped_annotations[filename], split_dir)
                for filename in files
            ]
            for future in futures:
                future.result()

    for split_name, files in splits.items():
        save_split(files, split_name)

json_file = '/content/traffic_light/train_dataset/train.json'
images_dir = '/content/traffic_light/train_dataset'
output_dir = 'yolo_dataset'
split_dataset(json_file, images_dir, output_dir)

Genrating yaml file

In [ ]:
data_yaml_content = """
train: yolo_dataset/train/images
val: yolo_dataset/val/images
test: yolo_dataset/test/images

nc: 2  # Number of classes (red, green)
names: ['red', 'green']  # Class names
"""

with open('data.yaml', 'w') as f:
    f.write(data_yaml_content)

Step 2: Training¶

In [ ]:
from ultralytics import YOLO

model = YOLO('yolov8s.pt')
model.train(data="/content/data.yaml", epochs=15, batch=16, imgsz=640)
Creating new Ultralytics Settings v0.0.6 file ✅ 
View Ultralytics Settings with 'yolo settings' or at '/root/.config/Ultralytics/settings.json'
Update Settings with 'yolo settings key=value', i.e. 'yolo settings runs_dir=path/to/dir'. For help see https://docs.ultralytics.com/quickstart/#ultralytics-settings.
Downloading https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8s.pt to 'yolov8s.pt'...
100%|██████████| 21.5M/21.5M [00:00<00:00, 156MB/s]
Ultralytics 8.3.49 🚀 Python-3.10.12 torch-2.5.1+cu121 CPU (Intel Xeon 2.20GHz)
engine/trainer: task=detect, mode=train, model=yolov8s.pt, data=/content/data.yaml, epochs=15, time=None, patience=100, batch=16, imgsz=640, save=True, save_period=-1, cache=False, device=None, workers=8, project=None, name=train, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, multi_scale=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, vid_stride=1, stream_buffer=False, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, embed=None, show=False, save_frames=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, show_boxes=True, line_width=None, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=True, opset=None, workspace=None, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, bgr=0.0, mosaic=1.0, mixup=0.0, copy_paste=0.0, copy_paste_mode=flip, auto_augment=randaugment, erasing=0.4, crop_fraction=1.0, cfg=None, tracker=botsort.yaml, save_dir=runs/detect/train
Downloading https://ultralytics.com/assets/Arial.ttf to '/root/.config/Ultralytics/Arial.ttf'...
100%|██████████| 755k/755k [00:00<00:00, 14.5MB/s]
Overriding model.yaml nc=80 with nc=2

                   from  n    params  module                                       arguments                     
  0                  -1  1       928  ultralytics.nn.modules.conv.Conv             [3, 32, 3, 2]                 
  1                  -1  1     18560  ultralytics.nn.modules.conv.Conv             [32, 64, 3, 2]                
  2                  -1  1     29056  ultralytics.nn.modules.block.C2f             [64, 64, 1, True]             
  3                  -1  1     73984  ultralytics.nn.modules.conv.Conv             [64, 128, 3, 2]               
  4                  -1  2    197632  ultralytics.nn.modules.block.C2f             [128, 128, 2, True]           
  5                  -1  1    295424  ultralytics.nn.modules.conv.Conv             [128, 256, 3, 2]              
  6                  -1  2    788480  ultralytics.nn.modules.block.C2f             [256, 256, 2, True]           
  7                  -1  1   1180672  ultralytics.nn.modules.conv.Conv             [256, 512, 3, 2]              
  8                  -1  1   1838080  ultralytics.nn.modules.block.C2f             [512, 512, 1, True]           
  9                  -1  1    656896  ultralytics.nn.modules.block.SPPF            [512, 512, 5]                 
 10                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 11             [-1, 6]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 12                  -1  1    591360  ultralytics.nn.modules.block.C2f             [768, 256, 1]                 
 13                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 14             [-1, 4]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 15                  -1  1    148224  ultralytics.nn.modules.block.C2f             [384, 128, 1]                 
 16                  -1  1    147712  ultralytics.nn.modules.conv.Conv             [128, 128, 3, 2]              
 17            [-1, 12]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 18                  -1  1    493056  ultralytics.nn.modules.block.C2f             [384, 256, 1]                 
 19                  -1  1    590336  ultralytics.nn.modules.conv.Conv             [256, 256, 3, 2]              
 20             [-1, 9]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 21                  -1  1   1969152  ultralytics.nn.modules.block.C2f             [768, 512, 1]                 
 22        [15, 18, 21]  1   2116822  ultralytics.nn.modules.head.Detect           [2, [128, 256, 512]]          
Model summary: 225 layers, 11,136,374 parameters, 11,136,358 gradients, 28.6 GFLOPs

Transferred 349/355 items from pretrained weights
TensorBoard: Start with 'tensorboard --logdir runs/detect/train', view at http://localhost:6006/
Freezing layer 'model.22.dfl.conv.weight'
train: Scanning /content/yolo_dataset/train/labels... 1819 images, 0 backgrounds, 0 corrupt: 100%|██████████| 1819/1819 [00:01<00:00, 1449.31it/s]
train: New cache created: /content/yolo_dataset/train/labels.cache

albumentations: Blur(p=0.01, blur_limit=(3, 7)), MedianBlur(p=0.01, blur_limit=(3, 7)), ToGray(p=0.01, num_output_channels=3, method='weighted_average'), CLAHE(p=0.01, clip_limit=(1.0, 4.0), tile_grid_size=(8, 8))
/usr/local/lib/python3.10/dist-packages/albumentations/__init__.py:24: UserWarning: A new version of Albumentations is available: 1.4.22 (you have 1.4.20). Upgrade using: pip install -U albumentations. To disable automatic update checks, set the environment variable NO_ALBUMENTATIONS_UPDATE to 1.
  check_for_updates()
val: Scanning /content/yolo_dataset/val/labels... 390 images, 0 backgrounds, 0 corrupt: 100%|██████████| 390/390 [00:00<00:00, 1535.20it/s]
val: New cache created: /content/yolo_dataset/val/labels.cache

Plotting labels to runs/detect/train/labels.jpg... 
optimizer: 'optimizer=auto' found, ignoring 'lr0=0.01' and 'momentum=0.937' and determining best 'optimizer', 'lr0' and 'momentum' automatically... 
optimizer: AdamW(lr=0.001667, momentum=0.9) with parameter groups 57 weight(decay=0.0), 64 weight(decay=0.0005), 63 bias(decay=0.0)
TensorBoard: model graph visualization added ✅
Image sizes 640 train, 640 val
Using 0 dataloader workers
Logging results to runs/detect/train
Starting training for 15 epochs...

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       1/15         0G      1.812      2.547     0.9416        149        640: 100%|██████████| 114/114 [1:19:27<00:00, 41.82s/it]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 13/13 [03:18<00:00, 15.26s/it]
                   all        390       2242      0.464      0.297      0.274      0.136

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       2/15         0G      1.861       1.46      0.937         83        640: 100%|██████████| 114/114 [1:19:09<00:00, 41.66s/it]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 13/13 [03:15<00:00, 15.02s/it]
                   all        390       2242       0.63      0.329      0.364      0.172

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       3/15         0G      1.873      1.418     0.9466         91        640: 100%|██████████| 114/114 [1:18:45<00:00, 41.45s/it]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 13/13 [03:14<00:00, 14.95s/it]
                   all        390       2242      0.567      0.334      0.342      0.147

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       4/15         0G      1.826      1.331     0.9431        109        640: 100%|██████████| 114/114 [1:17:51<00:00, 40.98s/it]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 13/13 [03:16<00:00, 15.14s/it]
                   all        390       2242      0.625      0.354      0.386      0.186

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       5/15         0G      1.743      1.214      0.918         72        640: 100%|██████████| 114/114 [1:18:51<00:00, 41.50s/it]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 13/13 [03:14<00:00, 14.96s/it]
                   all        390       2242      0.645      0.357      0.399       0.19

Closing dataloader mosaic
albumentations: Blur(p=0.01, blur_limit=(3, 7)), MedianBlur(p=0.01, blur_limit=(3, 7)), ToGray(p=0.01, num_output_channels=3, method='weighted_average'), CLAHE(p=0.01, clip_limit=(1.0, 4.0), tile_grid_size=(8, 8))

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       6/15         0G      1.756      1.242     0.9239         72        640: 100%|██████████| 114/114 [1:19:01<00:00, 41.59s/it]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 13/13 [03:15<00:00, 15.06s/it]
                   all        390       2242      0.591      0.363      0.381      0.185

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
       7/15         0G       1.77      1.222      0.895         60        640:   7%|â–‹         | 8/114 [05:38<1:14:20, 42.08s/it]

Step 3: Evaluation

In [ ]:
# Perform inference on test images
results = model.predict(source='yolo_dataset/test/images', save=True, conf=0.5)

# Display results
for result in results:
    print(result)
In [8]:
import pandas as pd
import matplotlib.pyplot as plt


results_file = "/content/runs/detect/train/results.csv"
data = pd.read_csv(results_file)


data['F1-Score'] = 2 * (data['metrics/precision(B)'] * data['metrics/recall(B)']) / (
    data['metrics/precision(B)'] + data['metrics/recall(B)']
)


plt.figure(figsize=(12, 6))
plt.plot(data['epoch'], data['train/box_loss'], label='Train Box Loss')
plt.plot(data['epoch'], data['train/cls_loss'], label='Train Classification Loss')
plt.plot(data['epoch'], data['train/dfl_loss'], label='Train DFL Loss')
plt.plot(data['epoch'], data['val/box_loss'], label='Validation Box Loss', linestyle='--')
plt.plot(data['epoch'], data['val/cls_loss'], label='Validation Classification Loss', linestyle='--')
plt.plot(data['epoch'], data['val/dfl_loss'], label='Validation DFL Loss', linestyle='--')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training and Validation Loss Curves')
plt.legend()
plt.grid(True)
plt.show()


plt.figure(figsize=(12, 6))
plt.plot(data['epoch'], data['metrics/precision(B)'], label='Precision (B)')
plt.plot(data['epoch'], data['metrics/recall(B)'], label='Recall (B)')
plt.plot(data['epoch'], data['metrics/mAP50(B)'], label='mAP@0.5 (B)')
plt.plot(data['epoch'], data['metrics/mAP50-95(B)'], label='mAP@0.5:0.95 (B)')
plt.xlabel('Epoch')
plt.ylabel('Metrics')
plt.title('Precision, Recall, and mAP Curves')
plt.legend()
plt.grid(True)
plt.show()


plt.figure(figsize=(12, 6))
plt.plot(data['epoch'], data['F1-Score'], label='F1-Score', color='blue')
plt.xlabel('Epoch')
plt.ylabel('F1-Score')
plt.title('F1-Score Curve')
plt.legend()
plt.grid(True)
plt.show()
No description has been provided for this image
No description has been provided for this image
No description has been provided for this image
In [12]:
def display_image(image_path):
    if not os.path.exists(image_path):
        print(f"File not found: {image_path}")
        return
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.figure(figsize=(10, 8))
    plt.imshow(image)
    plt.axis('off')
    plt.title(os.path.basename(image_path))
    plt.show()

train_batch_path = "/content/runs/detect/train/train_batch0.jpg"
display_image(train_batch_path)
No description has been provided for this image
In [14]:
import random
from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt
import os


model_path = "/content/runs/detect/train/weights/best.pt"
test_images_dir = "/content/yolo_dataset/train/images"


model = YOLO(model_path)


image_files = [os.path.join(test_images_dir, img) for img in os.listdir(test_images_dir) if img.endswith(('.png', '.jpg', '.jpeg'))]


random_images = random.sample(image_files, min(len(image_files), 5))

for image_path in random_images:
    results = model.predict(source=image_path, save=False, conf=0.25)


    for result in results:

        annotated_frame = result.plot()

        plt.figure(figsize=(10, 10))
        plt.imshow(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
        plt.title(f"Predictions for: {os.path.basename(image_path)}")
        plt.axis('off')
        plt.show()
image 1/1 /content/yolo_dataset/train/images/01065.jpg: 384x640 2 reds, 1 green, 386.0ms
Speed: 3.0ms preprocess, 386.0ms inference, 1.0ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image
image 1/1 /content/yolo_dataset/train/images/00547.jpg: 384x640 7 reds, 436.1ms
Speed: 4.4ms preprocess, 436.1ms inference, 1.3ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image
image 1/1 /content/yolo_dataset/train/images/02522.jpg: 384x640 1 red, 459.3ms
Speed: 4.3ms preprocess, 459.3ms inference, 8.4ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image
image 1/1 /content/yolo_dataset/train/images/02878.jpg: 384x640 1 red, 3 greens, 457.7ms
Speed: 4.3ms preprocess, 457.7ms inference, 1.0ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image
image 1/1 /content/yolo_dataset/train/images/01204.jpg: 384x640 1 green, 627.5ms
Speed: 4.9ms preprocess, 627.5ms inference, 1.4ms postprocess per image at shape (1, 3, 384, 640)
No description has been provided for this image